import os
import dotenv
import logging
from datetime import datetime
from typing import Dict, Any
from flask import Flask, request, jsonify
from flask_cors import CORS

# 加载环境变量
dotenv.load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 环境变量设置
os.environ["TAVILY_API_KEY"] = os.getenv("TAVILY_API_KEY")
os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("OPENAI_BASE_URL")



# 1. 工具调用优化
from langchain_community.tools.tavily_search import TavilySearchResults


def safe_tool_call(tool_func):
    def wrapper(*args, **kwargs):
        try:
            return tool_func(*args, **kwargs)
        except Exception as e:
            return f"工具调用失败: {str(e)}"

    return wrapper


# 创建搜索工具（增加结果数量）
search_tool = TavilySearchResults(max_results=3)

# 2. 文档检索系统
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools.retriever import create_retriever_tool


def setup_retrieval_system():
    """设置文档检索系统"""
    try:
        # 加载文档
        loader = WebBaseLoader("https://baike.baidu.com/item/%E7%8C%AB/22261")
        docs = loader.load()

        # 分割文档
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        documents = splitter.split_documents(docs)

        # 创建向量数据库
        embedding_model = OpenAIEmbeddings()
        vector = FAISS.from_documents(documents, embedding_model)

        # 创建搜索器检索策略
        retriever = vector.as_retriever(
            search_type="mmr",
            search_kwargs={"k": 3, "fetch_k": 10}
        )

        # 创建检索工具
        retriever_tool = create_retriever_tool(
            retriever=retriever,
            name="wiki_search",
            description="搜索维基百科关于猫的信息",
        )

        return retriever_tool
    except Exception as e:
        logger.error(f"检索系统设置失败: {e}")
        return None


# 3. 自定义提示词模板
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder


def create_custom_prompt():
    """创建自定义提示词"""
    return ChatPromptTemplate.from_messages([
        ("system", """你是一个有用的助手。请遵循以下规则：
1. 优先使用检索工具获取准确信息
2. 对于实时信息使用搜索工具
3. 保持回答简洁专业（100-500字）
4. 如果信息不足，明确说明
5. 记住对话历史中的用户信息"""),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad")
    ])


# 4. 创建记忆系统
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory


class EnhancedMemoryStore:
    def __init__(self):
        #self 的作用是指向当前类的实例对象，
        self.store = {}
        self.user_profiles = {}

    #取出当前用户记忆
    def get_session_history(self, session_id: str) -> BaseChatMessageHistory:
        if session_id not in self.store:
            #存入记忆
            self.store[session_id] = ChatMessageHistory()
            self.user_profiles[session_id] = {
                "created_at": datetime.now(),
                "interaction_count": 0,
                "user_name": None
            }
            logger.info(f"创建新会话: {session_id}")
        return self.store[session_id]

    #更新用户信息
    def update_user_profile(self, session_id: str, user_input: str):
        """更新用户档案"""
        if session_id in self.user_profiles:
            self.user_profiles[session_id]["interaction_count"] += 1

            # 简单提取用户名
            if "名字是" in user_input or "我叫" in user_input:
                if "名字是" in user_input:
                    #根据“名字是”进行分割，strip去除两端的空白字符
                    name_part = user_input.split("名字是")[1].strip()
                else:
                    name_part = user_input.split("我叫")[1].strip()

                name = name_part.split()[0] if name_part.split() else name_part
                self.user_profiles[session_id]["user_name"] = name
                logger.info(f"更新用户姓名: {name}")

    def get_user_info(self, session_id: str) -> Dict[str, Any]:
        """获取用户信息"""#get是字典对象里的方法，用于获取用于信息，找不到就返回默认值{}
        return self.user_profiles.get(session_id, {})


# 5. 响应验证器
from langchain_core.runnables import RunnableLambda


def create_response_validator():
    """创建响应验证器"""

    def validate_response(response):
        """验证和优化响应"""
        content = response.get("output", "")

        # 检查响应长度
        if len(content) > 1000:
            content = content[:1000] + "...(内容过长已截断)"
            logger.warning("响应过长已截断")

        # 简单的敏感信息检查
        sensitive_words = ["密码",  "API_KEY"]
        for word in sensitive_words:
            if word in content:
                content = "响应包含敏感信息，已过滤"
                logger.warning("检测到敏感信息")
                break

        response["output"] = content
        return response

    return RunnableLambda(validate_response)


# 6. 对话状态管理器,监控用户的情绪，状态
class ConversationManager:
    def __init__(self):
        self.conversation_states = {}

    def get_state(self, session_id: str) -> str:
        return self.conversation_states.get(session_id, "neutral")

    def update_state(self, session_id: str, user_input: str):
        """更新对话状态"""
        user_input_lower = user_input.lower()

        if any(word in user_input_lower for word in ["谢谢", "感谢", "thank"]):
            self.conversation_states[session_id] = "grateful"
        elif "?" in user_input or any(word in user_input_lower for word in ["什么", "怎么", "为什么"]):
            self.conversation_states[session_id] = "questioning"
        elif any(word in user_input_lower for word in ["你好", "嗨", "hello"]):
            self.conversation_states[session_id] = "greeting"
        else:
            self.conversation_states[session_id] = "neutral"


# 初始化全局变量
memory_store = EnhancedMemoryStore()#记忆系统对象
conversation_manager = ConversationManager()#对话状态管理对象
agent_with_chat_history = None

#初始化函数的构建，创建工具，执行器，智能体等，通过调用之前的函数
def initialize_agent():
    """初始化Agent"""
    global agent_with_chat_history

    try:
        # 设置文档的检索工具
        retriever_tool = setup_retrieval_system()
        if not retriever_tool:
            logger.error("检索工具设置失败，使用纯搜索模式")
            tools = [search_tool]
        else:
            tools = [search_tool, retriever_tool]

        # 创建模型
        from langchain_openai import ChatOpenAI
        model = ChatOpenAI(model="gpt-3.5-turbo")

        # 创建自定义提示词
        prompt = create_custom_prompt()

        # 创建Agent
        from langchain.agents import create_tool_calling_agent, AgentExecutor
        agent = create_tool_calling_agent(model, tools, prompt)

        # 创建执行器
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

        # 添加响应验证
        agent_with_validation = agent_executor | create_response_validator()

        # 创建带历史的Agent
        from langchain_core.runnables.history import RunnableWithMessageHistory
        agent_with_chat_history = RunnableWithMessageHistory(
            runnable=agent_with_validation,
            get_session_history=memory_store.get_session_history,
            input_messages_key="input",
            history_messages_key="chat_history",
        )

        logger.info("Agent初始化成功")
        return True
    except Exception as e:
        logger.error(f"Agent初始化失败: {e}")
        return False

# 创建Flask应用实例
app = Flask(__name__)
# 启用CORS并增加路由限制
CORS(app, resources={
    r"/api/*": {
        "origins": ["http://localhost:5173", "http://127.0.0.1:5173"],
        "methods": ["GET", "POST", "PUT", "DELETE"]
    }
})

# Flask路由，处理聊天的请求，获取结果，返还给前端
@app.route('/api/chat', methods=['POST'])
def chat():
    """处理聊天请求"""
    try:
        data = request.json
        user_input = data.get('message', '').strip()
        session_id = data.get('session_id', 'default_session')

        if not user_input:
            return jsonify({'error': '消息不能为空'}), 400

        if not agent_with_chat_history:
            return jsonify({'error': '聊天系统未初始化'}), 500

        # 更新对话状态和用户档案
        conversation_manager.update_state(session_id, user_input)
        memory_store.update_user_profile(session_id, user_input)

        # 获取响应
        response = agent_with_chat_history.invoke(
            {"input": user_input},
            config={"configurable": {"session_id": session_id}},
        )

        result = {
            'response': response.get('output', '无响应'),
            'session_id': session_id,
            'conversation_state': conversation_manager.get_state(session_id),
            'user_info': memory_store.get_user_info(session_id)
        }

        return jsonify(result)

    except Exception as e:
        logger.error(f"聊天处理错误: {e}")
        return jsonify({'error': '内部服务器错误'}), 500

#检查后端是否正常工作的路由
@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查"""
    return jsonify({'status': 'healthy', 'agent_initialized': agent_with_chat_history is not None})

#调用主函数
if __name__ == "__main__":
    # 初始化Agent
    if initialize_agent():
        logger.info("启动Flask服务器...")
        app.run(host='0.0.0.0', port=5000, debug=False)
    else:
        logger.error("Agent初始化失败，服务器无法启动")